# encoding:utf-8
# -----------------------------------------------------------
# "ML-Stealer: Stealing Prediction Functionality of Machine Learning Models with Mere Black-Box Access"
# @author: Shijie Wang, 2019.
# ------------------------------------------------------------

"""
data augmentation by VAE model
"""

from __future__ import print_function
from sklearn.externals import joblib
from vae_adult import VAE_adult
import torch
import random
import pandas as pd
import numpy as np


def adult_augmentation(vae_model, victim_model, numbers, outfolder):
    z = torch.randn(numbers, 3)
    output = vae_model.decode(z)
    output = output.detach().numpy()
    result = victim_model.predict_proba(output)
    index = np.where((result[:, 0] >= 0.95) | (result[:, 0] <= 0.2))
    output = output[index]

    target_columns = ['age', 'workclass', 'fnlwgt', 'education', 'education_number', 'marriage', 'occupation',
                      'relationship', 'race', 'sex', 'capital_gain', 'capital_loss', 'hours_per_work',
                      'native_country']

    adult_dict = {}
    for i in np.arange(0, 14):
        adult_dict[target_columns[i]] = output[:, i]

    adult_generation = pd.DataFrame(adult_dict)
    adult_generation.drop_duplicates(subset=None, keep='first', inplace=True)
    adult_generation.to_csv("%s/adult_no_svm_augmentation.csv" % outfolder, index=None)


def bank_augmentation(model, victim_model, numbers, outfolder):
    z = torch.randn(numbers, 3)
    output = model.decode(z)
    output = output.detach().numpy()
    result = victim_model.predict_proba(output)
    index = np.where((result[:, 0] >= 0.95) | (result[:, 0] <= 0.2))
    output = output[index]

    target_columns = ['age', 'job', 'marital', 'education', 'default', 'balance', 'housing', 'loan',
                      'contact', 'day', 'month', 'duration', 'campaign', 'pdays', 'previous', 'poutcome',
                      ]

    bank_dict = {}
    for i in np.arange(0, 16):
        bank_dict[target_columns[i]] = output[:, i]

    bank_generation = pd.DataFrame(bank_dict)
    bank_generation.drop_duplicates(subset=None, keep='first', inplace=True)
    bank_generation.to_csv("%s/bank_no_svm_augmentation.csv" % outfolder, index=None)


def cancer_augmentation(model, victim_model, numbers, outfolder):
    z = torch.randn(numbers, 3)
    output = model.decode(z)
    output = output.detach().numpy()
    result = victim_model.predict_proba(output)
    index = np.where((result[:, 0] >= 0.95) | (result[:, 0] <= 0.2))
    output = output[index]

    target_columns = ['code_number', 'thickness', 'cell_size', 'cell_shape', 'marginal_adhesion',
                      'single_cell_size', 'bare_nuclei', 'bland_chromatin', 'normal_nucleoli', 'mitoses'
                      ]

    cancer_dict = {}
    for i in np.arange(0, 10):
        cancer_dict[target_columns[i]] = output[:, i]

    cancer_generation = pd.DataFrame(cancer_dict)
    cancer_generation.drop_duplicates(subset=None, keep='first', inplace=True)
    cancer_generation.to_csv("%s/cancer_no_svm_augmentation.csv" % outfolder, index=None)


if __name__ == '__main__':
    manualSeed = random.randint(1, 10000)
    print("Random Seed: ", manualSeed)
    torch.manual_seed(manualSeed)
    device = torch.device('cpu')

    model_adult = torch.load('./vae_models/VAE_adult_svm_no_generator.pth', map_location='cpu')
    victim_model = joblib.load('./victim_models/svm_part_adult.m')
    adult_augmentation(model_adult, victim_model, 50000, './augmentation_data')
